#!/usr/bin/env python3
import time
import torch
import numpy as np
from logger import setup_logger
from yoloSeg import YoloSegPose

logger = setup_logger()

class ComputePosition:
    def __init__(self):
        self.yoloSegPose = YoloSegPose()
        
    def get_pose_dict(self, result, poly, threshold=0.8):
        """
        return {
            'left_top': {'coord': (x, y)},
            'left_bottom': {'coord': (x, y)},
            'right_top': {'coord': (x, y)},
            'left_0': {'coord': (x, y)},
            'top_0': {'coord': (x, y)},
            ... # 其他点
        }
        poly: 多边形轮廓点集，形状为(N,2)的数组
        """
        start_total = time.perf_counter()
        
        # 1. 确保多边形是二维数组
        start_poly = time.perf_counter()
        if poly.ndim == 3:
            poly = poly.reshape(-1, 2)
        poly_time = time.perf_counter() - start_poly
        logger.info(f"多边形数组处理耗时: {poly_time:.4f} 秒")

        # 2. 获取检测到的角点
        start_corners = time.perf_counter()
        corner_names = ["left_top", "left_bottom", "right_top", "right_bottom"]
        corner_indices = [0, 1, 2, 3]
        
        xy = result.xy
        conf = result.conf
        if conf is None:
            logger.warning("未检测到关键点")
            return None
        if isinstance(xy, torch.Tensor):
            xy = xy.cpu().numpy()
        if isinstance(conf, torch.Tensor):
            conf = conf.cpu().numpy()

        pose_dict = {}
        for name, idx in zip(corner_names, corner_indices):
            x, y = xy[0][idx]
            c = conf[0][idx]
            if c >= threshold:
                pose_dict[name] = {"coord": (float(x), float(y))}
        corners_time = time.perf_counter() - start_corners
        logger.info(f"获取角点耗时: {corners_time:.4f} 秒")
        
        # 3. 从多边形中移除已检测到的角点
        start_remove = time.perf_counter()
        remaining_poly = poly.copy()
        if remaining_poly.ndim == 3:
            remaining_poly = remaining_poly.reshape(-1, 2)
        for name in list(pose_dict.keys()):
            if name not in corner_names:
                continue
            cx, cy = pose_dict[name]["coord"]
            min_dist = float('inf')
            min_idx = -1
            for i, pt in enumerate(remaining_poly):
                px, py = pt
                dist = (px - cx) ** 2 + (py - cy) ** 2
                if dist < min_dist:
                    min_dist = dist
                    min_idx = i
            if min_idx >= 0:
                pose_dict[name] = {"coord": tuple(remaining_poly[min_idx])}
                remaining_poly = np.delete(remaining_poly, min_idx, axis=0)
        remove_time = time.perf_counter() - start_remove
        logger.info(f"移除已检测角点耗时: {remove_time:.4f} 秒")
        
        # 4. 推断边缘点
        start_edges = time.perf_counter()
        detected_corners = set(pose_dict.keys())
        if detected_corners == {"left_top","left_bottom","right_top","right_bottom"} and len(remaining_poly) == 0:
            logger.info("检测到情况0: 完整漏出四个点")
        elif detected_corners == {"left_top"} and len(remaining_poly) == 2:
            logger.info("检测到情况1: 只露出左上角，推断为三角形")
            remaining_poly = sorted(remaining_poly, key=lambda p: p[0])
            pose_dict["left_0"] = {"coord": tuple(remaining_poly[0])}
            pose_dict["top_0"] = {"coord": tuple(remaining_poly[1])}
        elif detected_corners == {"left_top", "left_bottom"} and len(remaining_poly) == 2:
            logger.info("检测到情况2: 露出左上+左下，推断为四边形")
            remaining_poly = sorted(remaining_poly, key=lambda p: p[1])
            pose_dict["top_0"] = {"coord": tuple(remaining_poly[0])}
            pose_dict["bottom_0"] = {"coord": tuple(remaining_poly[1])}
        elif detected_corners == {"left_top", "right_top"} and len(remaining_poly) == 2:
            logger.info("检测到情况3: 露出左上+右上，推断为四边形")
            remaining_poly = sorted(remaining_poly, key=lambda p: p[1])
            pose_dict["left_0"] = {"coord": tuple(remaining_poly[0])}
            pose_dict["right_0"] = {"coord": tuple(remaining_poly[1])}
        elif detected_corners == {"left_top", "left_bottom", "right_top"} and len(remaining_poly) == 2:
            logger.info("检测到情况4: 露出左上+左下+右上，推断为五边形")
            remaining_poly = sorted(remaining_poly, key=lambda p: p[1])
            pose_dict["right_0"] = {"coord": tuple(remaining_poly[0])}
            pose_dict["bottom_0"] = {"coord": tuple(remaining_poly[1])}
        elif detected_corners == {"left_bottom", "right_bottom"} and len(remaining_poly) == 2:
            logger.info("检测到情况5: 露出左下+右下，推断为四边形")
            remaining_poly = sorted(remaining_poly, key=lambda p: p[0])
            pose_dict["left_0"] = {"coord": tuple(remaining_poly[0])}
            pose_dict["right_0"] = {"coord": tuple(remaining_poly[1])}
        edges_time = time.perf_counter() - start_edges
        logger.info(f"推断边缘点耗时: {edges_time:.4f} 秒")
        
        total_time = time.perf_counter() - start_total
        logger.info(f"get_pose_dict 总耗时: {total_time:.4f} 秒")
        logger.info(f"pose_dict: {pose_dict}")
        return pose_dict
    
    def get_polygon_pose(self, image, seg_results, pose_results):
        start_total = time.perf_counter()
        
        fitted_pose = {}
        num_pose_points = 0
        if image is None:
            logger.warning("image is None")
            return num_pose_points, fitted_pose
        
        # 1. 关键点检测
        start_keypoints = time.perf_counter()
        if pose_results is None:
            logger.info("关键点为None")
            return num_pose_points, fitted_pose
        if len(pose_results.xy) == 0:
            return num_pose_points, fitted_pose
        keypoints_time = time.perf_counter() - start_keypoints
        logger.info(f"关键点检测耗时: {keypoints_time:.4f} 秒")
        
        # 2. 分割检测
        start_seg = time.perf_counter()
        if seg_results is None or len(seg_results) == 0:
            logger.warning("未检测到分割结果")
            return num_pose_points, fitted_pose
        seg_time = time.perf_counter() - start_seg
        logger.info(f"分割检测耗时: {seg_time:.4f} 秒")
        
        # 3. 获取原始多边形
        start_poly = time.perf_counter()
        polygon = seg_results
        poly_time = time.perf_counter() - start_poly
        logger.info(f"获取原始多边形耗时: {poly_time:.4f} 秒")
        
        # 4. 计算露出角点数量
        start_corners = time.perf_counter()
        xy = pose_results.xy
        conf = pose_results.conf
        if conf is None:
            num_pose_points = 0
        else:
            if isinstance(xy, torch.Tensor):
                xy = xy.cpu().numpy()
            if isinstance(conf, torch.Tensor):
                conf = conf.cpu().numpy()
            corner_confs = conf[0][:4]
            num_pose_points = np.sum(corner_confs >= 0.8)
        corners_time = time.perf_counter() - start_corners
        logger.info(f"计算角点数量耗时: {corners_time:.4f} 秒")
        
        # 5. 多边形拟合
        start_fit = time.perf_counter()
        fitted_polygon = self.yoloSegPose.get_polygon2(polygon, keypoints_num=num_pose_points)
        if fitted_polygon is None:
            logger.warning("多边形拟合失败")
            return num_pose_points, fitted_pose
        fit_time = time.perf_counter() - start_fit
        logger.info(f"多边形拟合耗时: {fit_time:.4f} 秒")
        
        # 6. 获取姿态点
        start_pose_dict = time.perf_counter()
        pose_dict = self.get_pose_dict(pose_results, fitted_polygon)
        pose_dict_time = time.perf_counter() - start_pose_dict
        logger.info(f"获取姿态点耗时: {pose_dict_time:.4f} 秒")
        
        if pose_dict is None:
            return num_pose_points, fitted_pose
        
        fitted_pose = {name: info["coord"] for name, info in pose_dict.items()}
        total_time = time.perf_counter() - start_total
        logger.info(f"get_polygon_pose 总耗时: {total_time:.4f} 秒")
        return num_pose_points, fitted_pose

    def get_edges_3d_points(self, color_img, depth_img, seg_results, pose_results, visual_pose=False,
                            visual_edges=False, visual_edges_3d=False, visual_fitted_edges=False, visual_fitted_vector=False):
        start_total = time.perf_counter()
        
        origin_img = color_img
        depth_img = depth_img
        
        # 1. 设置分辨率
        start_resolution = time.perf_counter()
        self.yoloSegPose.calibration_tool.set_resolution(origin_img, depth_img)
        resolution_time = time.perf_counter() - start_resolution
        logger.info(f"设置分辨率耗时: {resolution_time:.4f} 秒")
        
        # 2. 设置深度图
        start_depth = time.perf_counter()
        depth_img = self.yoloSegPose.calibration_tool.set_depth_img(depth_img)
        depth_time = time.perf_counter() - start_depth
        logger.info(f"设置深度图耗时: {depth_time:.4f} 秒")
        
        # 3. 获取多边形关键点
        start_pose = time.perf_counter()
        pose_num, pose = self.get_polygon_pose(origin_img, seg_results, pose_results)
        if pose_num == 0 or pose is None or len(pose) == 0:
            logger.warning("角点数量为零,不应该在这里出现才对")
            return None, None
        logger.info(f"获取图片多边形关键点{pose_num}个,角点{len(pose)}个")
        pose_time = time.perf_counter() - start_pose
        logger.info(f"获取多边形关键点耗时: {pose_time:.4f} 秒")
        
        if visual_pose:
            start_vis_pose = time.perf_counter()
            self.yoloSegPose.visualize_tool.visualize_pose(pose, origin_img)
            vis_pose_time = time.perf_counter() - start_vis_pose
            logger.info(f"可视化姿态耗时: {vis_pose_time:.4f} 秒")
        
        if len(pose) > 5:
            logger.warning("角点数量大于5,不可能出现才对")
            return None, None
        
        # 4. 获取边
        start_edges = time.perf_counter()
        edges = self.yoloSegPose.processPoint.get_dict_edges(pose, 60)
        if edges is not None and len(edges) > 0:
            logger.info(f"获取图片多边形边{len(edges)}条")
        else:
            logger.warning("未获取到多边形边,不应该在这里出现才对")
            return None, None
        edges_time = time.perf_counter() - start_edges
        logger.info(f"获取边耗时: {edges_time:.4f} 秒")
        
        # 5. 计算边深度
        start_edges_depth = time.perf_counter()
        edges_depth = self.yoloSegPose.processPoint.get_dict_depth(edges)
        logger.info(f"获取图片多边形边深度{len(edges_depth)}个")
        edges_depth_time = time.perf_counter() - start_edges_depth
        logger.info(f"计算边深度耗时: {edges_depth_time:.4f} 秒")
        
        if visual_edges:
            start_vis_edges = time.perf_counter()
            try:
                self.yoloSegPose.visualize_tool.visualize_edges(edges, edges_depth, origin_img, depth_img)
            except:
                logger.error("visual_edges error")
            vis_edges_time = time.perf_counter() - start_vis_edges
            logger.info(f"可视化边耗时: {vis_edges_time:.4f} 秒")
        
        # 6. 计算 3D 点
        start_3d_points = time.perf_counter()
        edges_3d = self.yoloSegPose.processPoint.get_dict_3dpoints(edges, edges_depth)
        logger.info(f"计算图片多边形边3D点{len(edges_3d)}个")
        edges_3d_time = time.perf_counter() - start_3d_points
        logger.info(f"计算3D点耗时: {edges_3d_time:.4f} 秒")
        
        if visual_edges_3d:
            start_vis_3d = time.perf_counter()
            try:
                self.yoloSegPose.visualize_tool.visualize_edges_3d(edges_3d)
            except:
                logger.error("visual_edges_3d error")
            vis_3d_time = time.perf_counter() - start_vis_3d
            logger.info(f"可视化3D边耗时: {vis_3d_time:.4f} 秒")
        
        # 7. 拟合点
        start_fit_points = time.perf_counter()
        fitted_edges = self.yoloSegPose.processLine.get_dict_fitted_points(edges_3d)
        logger.info(f"计算图片多边形边拟合点{len(fitted_edges)}个")
        fit_points_time = time.perf_counter() - start_fit_points
        logger.info(f"拟合点耗时: {fit_points_time:.4f} 秒")
        
        if visual_fitted_edges:
            start_vis_fit_edges = time.perf_counter()
            try:
                self.yoloSegPose.visualize_tool.visualize_fitted_edges(edges_3d, fitted_edges)
            except:
                logger.error("visual_fitted_edges error")
            vis_fit_edges_time = time.perf_counter() - start_vis_fit_edges
            logger.info(f"可视化拟合边耗时: {vis_fit_edges_time:.4f} 秒")
        
        # 8. 拟合向量
        start_fit_vector = time.perf_counter()
        fitted_vector = self.yoloSegPose.processLine.get_dict_fitted_vector(fitted_edges)
        logger.info(f"计算图片多边形边拟合向量成功")
        fit_vector_time = time.perf_counter() - start_fit_vector
        logger.info(f"拟合向量耗时: {fit_vector_time:.4f} 秒")
        
        if visual_fitted_vector:
            start_vis_vector = time.perf_counter()
            try:
                self.yoloSegPose.visualize_tool.visualize_fitted_vector(self.yoloSegPose.processVector.standard_vector, fitted_vector)
            except:
                logger.error("visual_fitted_vector error")
            vis_vector_time = time.perf_counter() - start_vis_vector
            logger.info(f"可视化拟合向量耗时: {vis_vector_time:.4f} 秒")
        
        total_time = time.perf_counter() - start_total
        logger.info(f"get_edges_3d_points 总耗时: {total_time:.4f} 秒")
        return pose, fitted_vector

    def test_one_image(self, color_img, depth_img, seg_results, pose_results):
        start_total = time.perf_counter()
        
        if color_img is None:
            logger.warning("color_img is None")
            return None, None
        if depth_img is None:
            logger.warning("depth_img is None")
            return None, None
        if seg_results is None:
            logger.warning("seg_results is None")
            return None, None
        if pose_results is None:
            logger.warning("pose_results is None")
            return None, None
        
        # 1. 获取 3D 点和拟合向量
        start_edges = time.perf_counter()
        pose, fitted_vector = self.get_edges_3d_points(color_img, depth_img, seg_results, pose_results)
        edges_time = time.perf_counter() - start_edges
        logger.info(f"获取3D点和拟合向量耗时: {edges_time:.4f} 秒")
        
        if pose is None:
            logger.warning("pose is None")
            return None, None
        if fitted_vector is None:
            logger.warning("fitted_vector is None")
            return None, None
        
        # 2. 计算向量
        start_vector = time.perf_counter()
        vector = self.yoloSegPose.processVector.get_xyzr_by_vector(fitted_vector)
        vector_time = time.perf_counter() - start_vector
        logger.info(f"计算向量耗时: {vector_time:.4f} 秒")
        
        if vector is None:
            logger.warning("vector is None")
            return None, None
        
        total_time = time.perf_counter() - start_total
        logger.info(f"test_one_image 总耗时: {total_time:.4f} 秒")
        return pose, vector

    def process_image_pair(self, timestamp, color_img, depth_img, seg_results, pose_results):
        start_total = time.perf_counter()
        
        pose, vector = self.test_one_image(color_img, depth_img, seg_results, pose_results)
        
        total_time = time.perf_counter() - start_total
        logger.info(f"process_image_pair 总耗时: {total_time:.4f} 秒")
        return pose, vector

    def set_standard(self, image_path, depth_path, seg_results, pose_results):
        start_total = time.perf_counter()
        
        # 1. 获取 3D 点和拟合向量
        start_edges = time.perf_counter()
        pose, fitted_vector = self.get_edges_3d_points(image_path, depth_path, seg_results, pose_results)
        edges_time = time.perf_counter() - start_edges
        logger.info(f"获取3D点和拟合向量耗时: {edges_time:.4f} 秒")
        
        if pose is not None and fitted_vector is not None:
            # 2. 设置标准向量
            start_standard = time.perf_counter()
            vector = self.yoloSegPose.processVector.set_standard(fitted_vector)
            standard_time = time.perf_counter() - start_standard
            logger.info(f"设置标准向量耗时: {standard_time:.4f} 秒")
            
            # 3. 计算向量
            start_vector = time.perf_counter()
            vector = self.yoloSegPose.processVector.get_xyzr_by_vector(fitted_vector)
            vector_time = time.perf_counter() - start_vector
            logger.info(f"计算向量耗时: {vector_time:.4f} 秒")
            
            total_time = time.perf_counter() - start_total
            logger.info(f"set_standard 总耗时: {total_time:.4f} 秒")
            return pose, vector
        else:
            total_time = time.perf_counter() - start_total
            logger.info(f"set_standard 总耗时: {total_time:.4f} 秒")
            return None, None

if __name__ == "__main__":
    compute_position = ComputePosition()